# knowledge/core/kb_application.py

import os
import logging
import tempfile
import shutil
import sys  # 添加这行
from typing import Dict, List, Optional, Any
from knowledge.database.db_manager import DatabaseManager
from knowledge.database.db_schema_manager import DBSchemaManager  # 添加Schema管理器

# 添加父目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
kb_root = os.path.dirname(os.path.dirname(current_dir))
sys.path.insert(0, kb_root)
from knowledge.util.tools import logging_helper

from knowledge.base.kb_base import KnowledgeBaseConfig
from knowledge.processors.kb_document_processor import KBDocumentProcessor
from knowledge.vector.kb_vector_store_manager import KBVectorStoreManager
from knowledge.query.kb_query_engine import KBQueryEngine

class KnowledgeBaseApplication:
    """主知识库应用类"""
    
    def __init__(self, chroma_db_path=None, collection_name="enterprise_knowledge", verbose=False):
        self.logger = logging.getLogger("KnowledgeBase.application")
        self.config = KnowledgeBaseConfig()
        self.doc_processor = KBDocumentProcessor()
        self.vector_manager = KBVectorStoreManager(chroma_db_path, collection_name, verbose=verbose)
        self.db_manager = DatabaseManager(verbose=verbose)  # 添加数据库管理器
        self.schema_manager = DBSchemaManager(verbose=verbose)  # 添加Schema管理器
        self.query_engine = None
        self.index = None
        self.knowledge_structure = {}
        self._engine_loaded = False
    
    def initialize(self, api_key=None):
        """初始化知识库系统"""
        try:
            # 1. 设置环境
            self.config.setup_environment()
            
            # 2. 初始化全局Settings配置
            self.config.initialize_global_settings(api_key)
            
            # 3. 加载文件注册表
            self.config.load_file_registry()
            
            logging_helper.info_if_verbose(self.logger,self.config.verbose,"知识库系统初始化完成")
            
            return True
            
        except Exception as e:
            self.logger.error(f"初始化失败: {e}")
            return False
    
    def ensure_query_engine_loaded(self):
        """确保查询引擎已加载 - 增强版"""
        if self._engine_loaded and self.query_engine:
            return True
            
        try:
            # 加载索引
            self.index, doc_count = self.vector_manager.create_or_load_index()
            if doc_count > 0:
                # 获取全局设置的LLM
                from llama_index.core import Settings
                llm = getattr(Settings, "llm", None)
                
                # 创建增强的查询引擎，传入数据库管理器和Schema管理器
                self.query_engine = KBQueryEngine(
                    self.index, 
                    db_manager=self.db_manager,
                    schema_manager=self.schema_manager,
                    llm=llm  # 传入LLM实例
                )
                self._engine_loaded = True
                logging_helper.info_if_verbose(self.logger,self.config.verbose,f"✅ 增强版查询引擎加载完成，文档数量: {doc_count}, LLM: {'已配置' if llm else '未配置'}")
                return True
            else:
                self.logger.error("❌ 没有可用的知识库索引")
                return False
        except Exception as e:
            self.logger.error(f"❌ 加载查询引擎失败: {e}")
            return False
    
    def build_knowledge_base(self, knowledge_root: str, force_rebuild: bool = False) -> Dict:
        """
        构建知识库
        
        Args:
            knowledge_root: 知识库根目录
            force_rebuild: 是否强制重建
            
        Returns:
            dict: 构建统计信息
        """
        if not self.config.is_initialized():
            raise Exception("请先调用initialize()方法初始化系统")
        
        # 检查目录是否存在
        if not os.path.exists(knowledge_root):
            self.logger.warning(f"知识库目录不存在: {knowledge_root}")
            return {
                "total_files": 0,
                "updated_files": 0,
                "total_documents": 0,
                "folders": [],
                "error": f"目录不存在: {knowledge_root}"
            }
        
        # 首先设置存储上下文，检查集合状态
        storage_context, chroma_collection = self.vector_manager.setup_storage()
        collection_exists = chroma_collection.count() > 0
        
        logging_helper.info_if_verbose(self.logger,self.config.verbose,f"集合状态: 存在={collection_exists}, 文档数量={chroma_collection.count()}")
        
        # 如果集合存在且不是强制重建，直接加载索引
        if collection_exists and not force_rebuild:
            logging_helper.info_if_verbose(self.logger,self.config.verbose,f"找到现有向量数据（{chroma_collection.count()} 个文档），直接加载索引...")
            self.index, doc_count = self.vector_manager.create_or_load_index()
            self.query_engine = KBQueryEngine(self.index)
            
            stats = {
                "total_files": "无需扫描",  # 不需要扫描文件
                "updated_files": 0,
                "total_documents": doc_count,
                "folders": "从现有索引加载",
                "loaded_from_cache": True
            }
            
            logging_helper.info_if_verbose(self.logger,self.config.verbose,f"索引加载完成，文档数量: {doc_count}")
            return stats
        
        # 需要构建或重建索引的情况
        logging_helper.info_if_verbose(self.logger,self.config.verbose,f"{'强制重建' if force_rebuild else '构建新'}知识库索引...")
        
        # 1. 扫描知识库结构
        logging_helper.info_if_verbose(self.logger,self.config.verbose,f"扫描知识库目录: {knowledge_root}")
        self.knowledge_structure = self.doc_processor.scan_knowledge_base(knowledge_root)
        
        # 2. 收集所有文件路径
        all_file_paths = []
        for folder, files in self.knowledge_structure.items():
            all_file_paths.extend(files)
            logging_helper.info_if_verbose(self.logger,self.config.verbose,f"文件夹 '{folder}': {len(files)} 个文件")
        
        if not all_file_paths:
            self.logger.warning("未找到任何文档文件")
            return {
                "total_files": 0,
                "updated_files": 0,
                "total_documents": 0,
                "folders": list(self.knowledge_structure.keys()),
                "error": "未找到任何文档文件"
            }
        
        # 3. 加载文档
        documents = self.doc_processor.load_documents_with_metadata(all_file_paths, knowledge_root)
        
        # 4. 构建或更新索引
        if force_rebuild:
            logging_helper.info_if_verbose(self.logger,self.config.verbose,"强制重建知识库索引...")
            self.index, doc_count = self.vector_manager.create_or_load_index(
                documents=documents,
                force_rebuild=True
            )
            updated_count = len(documents)
        else:
            logging_helper.info_if_verbose(self.logger,self.config.verbose,"增量更新知识库索引...")
            updated_count = self.vector_manager.incremental_update_index(
                documents=documents,
                file_registry=self.config.file_registry,
                config=self.config
            )
            self.index, doc_count = self.vector_manager.create_or_load_index()
        
        # 5. 创建查询引擎
        if doc_count > 0:
            self.query_engine = KBQueryEngine(self.index)
        
        # 6. 保存文件注册表
        self.config.save_file_registry()
        
        stats = {
            "total_files": len(all_file_paths),
            "updated_files": updated_count,
            "total_documents": doc_count,
            "folders": list(self.knowledge_structure.keys()),
            "loaded_from_cache": False
        }
        
        logging_helper.info_if_verbose(self.logger,self.config.verbose,f"知识库构建完成: {stats}")
        return stats
    
    def query(self, query_text: str, filters: Optional[Dict] = None, show_details: bool = True) -> Dict:
        """执行查询 - 现在使用智能融合查询"""
        if not self.ensure_query_engine_loaded():
            raise Exception("查询引擎未初始化")
        
        # 使用新的智能融合查询
        return self.query_engine.integrated_query(query_text, show_details)
    
    def batch_query(self, queries: List[str], filters: Optional[Dict] = None, show_details: bool = True) -> List[Dict]:
        """批量查询"""
        if not self.query_engine:
            raise Exception("请先调用build_knowledge_base()方法构建知识库")
        
        return [self.query(q, filters, show_details) for q in queries]
    
    def get_knowledge_structure(self) -> Dict:
        """获取知识库结构"""
        return self.knowledge_structure
    
    def get_collection_info(self) -> Dict:
        """获取集合信息"""
        return self.vector_manager.get_collection_stats()
    
    def search_by_folder(self, query_text: str, folder_name: str, show_details: bool = True) -> Dict:
        """按文件夹搜索"""
        filters = {"folder": folder_name}
        return self.query(query_text, filters, show_details)
    
    def is_index_loaded(self) -> bool:
        """检查索引是否已加载"""
        return self.query_engine is not None
    
    def smart_query(self, query_text: str, use_database: bool = True, show_details: bool = True) -> Dict:
        """
        智能查询：优先尝试数据库查询，失败则回退到知识库查询
        
        Args:
            query_text: 查询文本
            use_database: 是否尝试数据库查询
            show_details: 是否显示详情
        """
        if not self.ensure_query_engine_loaded():
            raise Exception("查询引擎未初始化")
        
        # 判断是否可能是数据库查询
        if use_database and self._is_likely_database_query(query_text):
            logging_helper.info_if_verbose(self.logger,self.config.verbose,"检测到可能的数据库查询，尝试智能数据库查询")
            return self.query_engine.query_with_database(query_text, show_details)
        else:
            logging_helper.info_if_verbose(self.logger,self.config.verbose,"使用标准知识库查询")
            # 修复：只传递 2-3 个参数，将 show_details 作为关键字参数或调整调用方式
            return self.query_engine.query(query_text, show_details=show_details)
    
    def _is_likely_database_query(self, query_text: str) -> bool:
        """判断查询是否可能是数据库相关的"""
        db_keywords = [
            '项目', '商机', '员工', '产品', '金额', '数量', '统计', 
            '多少', '总计', '汇总', '查询', '查找', '搜索',
            'project', 'pipline', 'employee', 'product', 'amount', 'count'
        ]
        query_lower = query_text.lower()
        return any(keyword in query_lower for keyword in db_keywords)
    
    def integrated_query(self, query_text: str, show_details: bool = True) -> Dict:
        """
        智能融合查询：结合数据库查询和知识库查询
        
        Args:
            query_text: 查询文本
            show_details: 是否显示详情
            
        Returns:
            dict: 查询结果
        """
        if not self.ensure_query_engine_loaded():
            raise Exception("查询引擎未初始化")
        
        return self.query_engine.integrated_query(query_text, show_details)